package com.demo.config;

import com.demo.filter.RequestContext;
import com.demo.filter.RequestContextHelper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.validation.constraints.NotNull;
import java.util.Map;

/**
 * @author wangfengchen
 * 上下文传递装饰器
 */
@Slf4j
public class ContextAwareTaskDecorator implements TaskDecorator {
    @NotNull
    @Override
    public Runnable decorate(@NotNull Runnable runnable) {
        ServletRequestAttributes context = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        RequestContext requestContext = RequestContextHelper.get();
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                if (ObjectUtils.isNotEmpty(contextMap)) {
                    MDC.setContextMap(contextMap);
                }
                if (ObjectUtils.isNotEmpty(requestContext)) {
                    RequestContextHelper.set(requestContext);
                }
                if (ObjectUtils.isNotEmpty(context)) {
                    RequestContextHolder.setRequestAttributes(context);
                }
                runnable.run();
            } finally {
                RequestContextHelper.remove();
                MDC.clear();
                try {
                    RequestContextHolder.resetRequestAttributes();
                } catch (Exception e) {
                    log.info("清理线程依赖信息异常->{},{}", e.getMessage(), ExceptionUtils.getStackTrace(e));
                }
            }
        };
    }
}